# -*- coding: utf-8 -*-
"""
Created on Fri Sep  6 14:09:54 2019

@author: Train a link prediction model based on GCN embeddings.
"""
import pandas as pd
import numpy as np
import json
import random
import gensim
from gensim.models import KeyedVectors
from sklearn.model_selection import train_test_split
import keras
from keras.layers import Dense,Dropout
from keras.models import Sequential

with open('../data/friends_dict.json', 'r') as json_file:
    json_str = json_file.read()
friends_dict = json.loads(json_str) # 这里读进来的键是str类型，和值是一个列表，里面是int
relation_df = pd.read_csv('../data/true_friends.csv')
gs = KeyedVectors.load_word2vec_format("../data/gs_data/emb/pokec_gcn_feat.emb")


def link_vec(node1,node2,operator):
    assert operator == 'average' or operator == 'hadamard', \
    "Unknown operator! you can only choose 'average' or 'hadamard'"
    v1,v2 = gs[str(node1)],gs[str(node2)]
    if operator == 'average':
        return (v1+v2)/2
    if operator == 'hadamard':
        return v1*v2

all_nodes = list(set(relation_df.id1))
nodes = all_nodes
pos_pairs = []
neg_pairs = []
for node in nodes:
    for friend in friends_dict[str(node)]:
        pos_pairs.append((node,friend))
        neg_pairs.append((node,random.choice(all_nodes)))
        
        
print("Num of positive pairs:",len(pos_pairs))
print("Num of negative pairs:",len(neg_pairs))

X_pos = [link_vec(pair[0],pair[1],'hadamard') for pair in pos_pairs]
y_pos = [1]*len(X_pos)
X_neg = [link_vec(pair[0],pair[1],'hadamard') for pair in neg_pairs]
y_neg = [0]*len(X_neg)
X = X_pos + X_neg
y = y_pos + y_neg

X_train, X_test, y_train, y_test = train_test_split(np.array(X), np.array(y), test_size=0.33, random_state=42)


m = Sequential()
m.add(Dense(128,activation='relu'))
m.add(Dense(64,activation='relu'))
m.add(Dense(1,activation='sigmoid'))
m.compile(optimizer='adam',loss='binary_crossentropy',metrics=['accuracy'])

m.fit(X_train,y_train,batch_size=64,epochs=10,validation_split=0.15)

m.evaluate(X_test,y_test)

m.save('../match_models/LM_gcn_feat1.h5')

def link_predict(node1,node2):
    lv = link_vec(node1,node2,'hadamard')
    predict = m.predict(np.array([lv]))[0][0]
    if predict>0.5:
        return 1
    else:
        return 0

def evaluate(node):
    cnt = 0
    ts = friends_dict[str(1048586)]
    fs = friends_dict[str(node)]
    for t in ts:
        if t in fs:
            true = 1
        else:
            true = 0
        if link_predict(node,t) == true:
            cnt += 1
    return cnt/len(ts)

acc = []
for node in nodes[:100]:
    acc.append(evaluate(node))
    print(evaluate(node))
print('Average:',sum(acc)/len(acc))




